import os
import shutil
import cv2
import json
from tqdm import tqdm

# Set input and output directories
INPUT_DIR = "images/img"
OUTPUT_DIR = "processed_images"
OUTPUT_IMG_DIR = os.path.join(OUTPUT_DIR, "img")  # 新增图片输出目录
INPUT_JSON = "images/data.json"
OUTPUT_JSON = "processed_images/data.json"

def is_valid_image(file_path):
    # Common image file signatures (magic bytes)
    signatures = {
        b"\xff\xd8\xff": "JPEG",
        b"\x89PNG\r\n\x1a\n": "PNG",
        b"GIF87a": "GIF",
        b"GIF89a": "GIF",
        b"BM": "BMP",
        b"II*\x00": "TIFF",
        b"MM\x00*": "TIFF",
    }

    with open(file_path, "rb") as f:
        header = f.read(8)  # Read first 8 bytes

    for sig, fmt in signatures.items():
        if header.startswith(sig):
            return True
    return False

def extract_first_frame(video_path, output_path):
    # Open the video file
    video = cv2.VideoCapture(video_path)

    # Check if video opened successfully
    if not video.isOpened():
        raise Exception(f"Could not open video {video_path}")

    # Read the first frame
    ret, frame = video.read()

    if not ret:
        raise Exception(f"Could not read frame from video {video_path}")

    # Release the video capture object
    video.release()
    
    return frame

def process_image(image):
    """
    处理图片的函数，可以根据需求添加更多处理步骤
    """
    # 确保图片尺寸合适（例如限制最大尺寸为1024像素）
    max_size = 1024
    height, width = image.shape[:2]
    
    if height > max_size or width > max_size:
        if height > width:
            new_height = max_size
            new_width = int(width * (max_size / height))
        else:
            new_width = max_size
            new_height = int(height * (max_size / width))
            
        image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
    
    return image

def main():
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    os.makedirs(OUTPUT_IMG_DIR, exist_ok=True)  # 创建图片输出目录
    
    # 读取原始 JSON 数据
    print("Loading data.json...")
    with open(INPUT_JSON, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 获取所有输入文件
    files = os.listdir(INPUT_DIR)
    
    print(f"Found {len(files)} files in {INPUT_DIR}")
    print("Processing files...")
    
    success_count = 0
    error_count = 0
    processed_data = {}
    failed_files = set()
    
    for filename in tqdm(files):
        input_path = os.path.join(INPUT_DIR, filename)
        output_path = os.path.join(OUTPUT_IMG_DIR, filename)  # 修改输出路径到 img 子目录
        
        try:
            if is_valid_image(input_path):
                # 如果是有效的图片文件，直接读取并处理
                image = cv2.imread(input_path)
                if image is None:
                    raise Exception("Failed to read image")
                
                processed_image = process_image(image)
                cv2.imwrite(output_path, processed_image)
                success_count += 1
                
            else:
                # 尝试作为视频处理
                try:
                    frame = extract_first_frame(input_path, output_path)
                    processed_frame = process_image(frame)
                    cv2.imwrite(output_path, processed_frame)
                    success_count += 1
                except Exception as e:
                    raise Exception(f"Not a valid image or video file: {e}")
                    
        except Exception as e:
            print(f"Error processing {filename}: {e}")
            error_count += 1
            failed_files.add(filename)
            continue
    
    # 处理 JSON 数据
    print("\nUpdating data.json...")
    for item_id, item in data.items():
        # 获取文件名（处理 Windows 风格的路径）
        filename = os.path.basename(item["url"].replace("\\", "/"))
        
        # 如果图片处理成功，保留数据并更新 url 路径格式
        if filename not in failed_files:
            # 更新 url 字段，统一使用正斜杠
            item["url"] = item["url"].replace("\\", "/")
            processed_data[item_id] = item
    
    # 保存更新后的 JSON 数据
    with open(OUTPUT_JSON, 'w', encoding='utf-8') as f:
        json.dump(processed_data, f, ensure_ascii=False, indent=2)
    
    print("\nProcessing complete!")
    print(f"Successfully processed: {success_count} files")
    print(f"Failed to process: {error_count} files")
    print(f"Original JSON entries: {len(data)}")
    print(f"Remaining JSON entries: {len(processed_data)}")
    print(f"Processed images saved to: {OUTPUT_DIR}")
    print(f"Updated data.json saved to: {OUTPUT_JSON}")

if __name__ == "__main__":
    main() 